from .layers import *

import os
import numpy as np
import tensorflow as tf

initializer = tf.contrib.layers.xavier_initializer()


def isensee2017_model(inputs, training, keep_prob):

    n_base_filters = 16
    depth = 3
    n_segmentation_levels = 2

    # inputs = tf.concat([image, memory], -1)

    current_layer = inputs
    level_output_layers = list()
    level_filters = list()

    for level_number in range(depth):
        n_level_filters = (2**level_number) * n_base_filters
        level_filters.append(n_level_filters)

        if current_layer is inputs:
            in_conv = Conv(current_layer, n_level_filters, 3, 1, keep_prob, 'cil')
        else:
            in_conv = Conv(current_layer, n_level_filters, 3, 2, keep_prob, 'cil')

        summation_layer = Conv(in_conv, n_level_filters, 3, 1, keep_prob, 'cilscila')
        level_output_layers.append(summation_layer)
        current_layer = summation_layer

    segmentation_layers = list()
    for level_number in range(depth - 2, -1, -1):
        n_filters = level_filters[level_number]
        up_sampling = Conv(tf.keras.layers.UpSampling3D(2)(current_layer), n_filters, 3, 1, keep_prob, 'cil')
        concatenation_layer = tf.concat([level_output_layers[level_number], up_sampling], -1)
        localization_output = Conv(concatenation_layer, n_filters, 3, 1, keep_prob, 'cil')
        localization_output = Conv(localization_output, n_filters, 1, 1, keep_prob, 'cil')
        current_layer = localization_output
        if level_number < n_segmentation_levels:
            segmentation_layers.insert(0, Conv(current_layer, 1, 1, 1, keep_prob, 'c'))

    output_layer = None
    for level_number in reversed(range(n_segmentation_levels)):
        segmentation_layer = segmentation_layers[level_number]
        if output_layer is None:
            output_layer = segmentation_layer
        else:
            output_layer = output_layer + segmentation_layer

        if level_number > 0:
            output_layer = tf.keras.layers.UpSampling3D(size=(2, 2, 2))(output_layer)

    return output_layer